/*
 * *
 *  * blog.coder4j.cn
 *  * Copyright (C) 2016-2019 All Rights Reserved.
 *
 */
package com.chen.common.websocket.handler;

import cn.hutool.core.lang.Assert;
import com.chen.common.websocket.annotation.WS;
import com.chen.common.websocket.utils.WSUitl;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.*;

/**
 * @author ChenJinLong
 */
@Slf4j
@Component
public class WsSessionHandler extends TextWebSocketHandler {

    /**
     * socket 建立成功事件
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        // 添加 session
        getWsHandler(session).sessionMap.put(WSUitl.key(session), session);

        log.info("【WebSocket】 -> 连接完毕, 总连接数: {}", WSUitl.totalConnectNum());
    }

    /**
     * 接收消息事件
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        getWsHandler(session).handleTextMessage(message);
    }

    /**
     * socket 断开连接时
     */
    @Override
    @SneakyThrows
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        AbsWebSocketHandler wsHandler = getWsHandler(session);
        WebSocketSession wsSession = wsHandler.sessionMap.remove(WSUitl.key(session));
        wsSession.close();
        log.info("【WebSocket】 -> 连接断开, 总连接数: {}", WSUitl.totalConnectNum());
    }

    private AbsWebSocketHandler getWsHandler(WebSocketSession session) {
        Map<String, AbsWebSocketHandler> wsHandlerMap = WSUitl.wsHandlerMap();

        AbsWebSocketHandler absWsHandler = null;
        List<String> extremityList = new ArrayList<>();
        for (AbsWebSocketHandler handler : wsHandlerMap.values()) {
            WS ws = handler.getClass().getAnnotation(WS.class);

            if (ws != null){
                extremityList.add(ws.value());
                if (ws.value().equals(WSUitl.extremity(session))){
                    absWsHandler = handler;
                }
            }
        }

        Set<String> extremitySet = new HashSet<>(extremityList);
        Assert.isTrue(extremitySet.size() == extremityList.size(), "@WS 端点重复");
        Assert.notNull(absWsHandler, "【WebSocket】AbsWebSocketHandler不存在, " +
                "请检查AbsWebSocketHandler的实现类是否有@WS注解, 或者@WS的端点填写是否正确");
        return absWsHandler;
    }
}